from mindformers import CLIPModel, CLIPProcessor
from mindspore import Tensor
import mindspore as ms
ms.set_context(mode=0, device_id=7, device_target="CPU")

from PIL import Image
import requests
from transformers import AutoProcessor, CLIPVisionModel

processor = AutoProcessor.from_pretrained("/home/zhangyouwen/suite/mobile_commucation/clip")

url = "http://images.cocodataset.org/val2017/000000039769.jpg"
image = Image.open(requests.get(url, stream=True).raw)

inputs = Tensor(processor(images=image, return_tensors="np").pixel_values)
model = CLIPModel.from_pretrained('clip_vit_l_14@336')
output = model.get_image_features(inputs).asnumpy()

model = CLIPVisionModel.from_pretrained("/home/zhangyouwen/suite/mobile_commucation/clip")
inputs = processor(images=image, return_tensors="pt")

outputs = model(**inputs)
last_hidden_state = outputs.last_hidden_state.detach().numpy()

shape = output.shape
assert output.shape == last_hidden_state.shape


diff_res = [0, 0, 0, 0]

for i in range(shape[0]):
    for j in range(shape[1]):
        for k in range(shape[2]):
            diff = abs(output[i][j][k] - last_hidden_state[i][j][k])
            if diff * 10 >= 1:
                diff_res[0] += 1
            if diff * 100 >= 1:
                diff_res[1] += 1
            if diff * 1000 >= 1:
                diff_res[2] += 1
            if diff * 10000 >= 1:
                diff_res[3] += 1

num = shape[0] * shape[1] * shape[2]
diff_res = [i / num for i in diff_res]
# diff_res = [0.0, 0.0, 0.0011034987001733102, 0.01511725519930676]
print(f"logits的十分位不同的占比为{diff_res[0]}")
print(f"logits的百分位不同的占比为{diff_res[1]}")
print(f"logits的千分位不同的占比为{diff_res[2]}")
print(f"logits的万分位不同的占比为{diff_res[3]}")
print(diff_res)